package com.raos.example.service;

import com.fasterxml.jackson.databind.JsonNode;
import com.raos.example.utils.JsonUtil;
import com.raos.example.utils.OllamaUtil;
import lombok.extern.slf4j.Slf4j;
import org.noear.solon.annotation.Component;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Consumer;

/**
 * Ollama 连接服务层
 *
 * @author raos
 * @date 2025-02-14
 */
@Slf4j
@Component
public class OllamaService {

    private static final String OLLAMA_URL = OllamaUtil.OLLAMA_API_URL;
    private static final String DEFAULT_DEEPSEEK_MODEL_NAME = OllamaUtil.DEFAULT_DEEPSEEK_MODEL_NAME;

    /**
     * 创建对话响应信息
     *
     * @param modelName 模型名
     * @param prompt 提示词
     * @param callback 回调响应信息
     */
    public void generateResponse(String modelName, String prompt, Consumer<String> callback) {
        Map<String, Object> data = new HashMap<>(4);
        data.put("model", modelName);
        data.put("prompt", prompt.replace("\"", "\\\""));
        data.put("stream", true);
        String requestJsonStr = JsonUtil.toJsonStr(data);
        System.out.println(requestJsonStr);

        String requestBody = String.format("""
        {
            "model": "%s",
            "prompt": "%s",
            "stream": true
        }
        """, modelName, prompt.replace("\"", "\\\""));

        HttpClient client = HttpClient.newHttpClient();
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(OLLAMA_URL))
                .header("Content-Type", "application/json")
                .POST(HttpRequest.BodyPublishers.ofString(requestBody))
                .build();

        client.sendAsync(request, HttpResponse.BodyHandlers.ofLines())
                .thenAccept(response -> {
                    response.body().forEach(line -> {
                        log.debug("line：[{}]", line);
                        JsonNode json = JsonUtil.toNode(line);
                        if (json != null && json.has("response")) {
                            callback.accept(json.get("response").asText());
                        }
                    });
                });
    }

}
